//! 增强功能演示程序
//!
//! 展示新实现的核心功能：
//! 1. Flight SQL准备语句支持
//! 2. Flight SQL事务管理
//! 3. HTTP网关认证授权
//! 4. S3对象存储数据源

use std::collections::HashMap;
use std::sync::Arc;
use tokio;
use data_gateway_kernel::DataGatewayKernel;
use data_gateway_flight_sql::service_simple::SimpleFlightSqlService;
use data_gateway_gateway::middleware::{MemoryAuthProvider, AuthMethod, MiddlewareManager, AuthProvider, RateLimitConfig};
use data_gateway_data_sources::{S3DataSource, S3Config, DataSourceQuery, DataSourceManager, DataSource, ElasticsearchDataSource, ElasticsearchConfig};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    println!("🚀 数据网关增强功能演示程序");
    println!("========================================");

    // 1. 演示Flight SQL准备语句功能
    println!("\n📋 1. Flight SQL准备语句功能演示");
    demo_prepared_statements().await?;

    // 2. 演示Flight SQL事务管理功能
    println!("\n💼 2. Flight SQL事务管理功能演示");
    demo_transaction_management().await?;

    // 3. 演示HTTP网关认证授权功能
    println!("\n🔐 3. HTTP网关认证授权功能演示");
    demo_authentication_authorization().await?;

    // 4. 演示S3对象存储数据源功能
    println!("\n☁️ 4. S3对象存储数据源功能演示");
    demo_s3_data_source().await?;

    // 5. 演示Elasticsearch搜索引擎数据源功能
    println!("\n🔍 5. Elasticsearch搜索引擎数据源功能演示");
    demo_elasticsearch_data_source().await?;

    // 6. 演示限流控制中间件功能
    println!("\n🚦 6. 限流控制中间件功能演示");
    demo_rate_limiting().await?;

    // 7. 综合功能演示
    println!("\n🎯 7. 综合功能演示");
    demo_integrated_features().await?;

    println!("\n✅ 所有增强功能演示完成！");
    println!("========================================");

    Ok(())
}

/// 演示Flight SQL准备语句功能
async fn demo_prepared_statements() -> anyhow::Result<()> {
    println!("  创建Flight SQL服务...");
    let kernel = Arc::new(DataGatewayKernel::new().await?);
    let flight_service = SimpleFlightSqlService::new(kernel);

    println!("  📝 创建准备语句...");
    let sql = "SELECT * FROM users WHERE id = ? AND status = ?";
    let stmt_id = flight_service.create_prepared_statement(sql).await?;
    println!("    ✅ 准备语句创建成功，ID: {}", stmt_id);

    println!("  📊 获取准备语句信息...");
    let stmt_info = flight_service.get_prepared_statement(&stmt_id).await?;
    println!("    ✅ 语句SQL: {}", stmt_info.sql);
    println!("    ✅ 创建时间: {:?}", stmt_info.created_at);

    println!("  🔄 执行准备语句...");
    let parameters = HashMap::new(); // 简化实现，暂不处理参数
    match flight_service.execute_prepared_statement(&stmt_id, parameters).await {
        Ok(_) => println!("    ✅ 准备语句执行成功"),
        Err(e) => println!("    ⚠️ 准备语句执行失败（预期，因为简化实现）: {}", e),
    }

    println!("  🗑️ 关闭准备语句...");
    flight_service.close_prepared_statement(&stmt_id).await?;
    println!("    ✅ 准备语句关闭成功");

    Ok(())
}

/// 演示Flight SQL事务管理功能
async fn demo_transaction_management() -> anyhow::Result<()> {
    println!("  创建Flight SQL服务...");
    let kernel = Arc::new(DataGatewayKernel::new().await?);
    let flight_service = SimpleFlightSqlService::new(kernel);

    println!("  🚀 开始事务...");
    let tx_id = flight_service.begin_transaction().await?;
    println!("    ✅ 事务创建成功，ID: {}", tx_id);

    println!("  📊 获取事务信息...");
    let tx_info = flight_service.get_transaction(&tx_id).await?;
    println!("    ✅ 事务状态: {:?}", tx_info.status);
    println!("    ✅ 创建时间: {:?}", tx_info.created_at);

    println!("  💾 提交事务...");
    flight_service.commit_transaction(&tx_id).await?;
    println!("    ✅ 事务提交成功");

    println!("  📊 验证事务状态...");
    let tx_info = flight_service.get_transaction(&tx_id).await?;
    println!("    ✅ 事务状态: {:?}", tx_info.status);

    // 演示回滚功能
    println!("  🔄 演示事务回滚...");
    let tx_id2 = flight_service.begin_transaction().await?;
    println!("    ✅ 新事务创建成功，ID: {}", tx_id2);
    
    flight_service.rollback_transaction(&tx_id2).await?;
    println!("    ✅ 事务回滚成功");

    let tx_info2 = flight_service.get_transaction(&tx_id2).await?;
    println!("    ✅ 回滚后事务状态: {:?}", tx_info2.status);

    Ok(())
}

/// 演示HTTP网关认证授权功能
async fn demo_authentication_authorization() -> anyhow::Result<()> {
    println!("  创建认证提供者...");
    let auth_provider = MemoryAuthProvider::new();

    println!("  🔑 测试API密钥认证...");
    let api_key_method = AuthMethod::ApiKey("admin-key-123".to_string());
    let auth_result = auth_provider.authenticate(api_key_method).await;
    match auth_result {
        data_gateway_gateway::middleware::AuthResult::Success(user) => {
            println!("    ✅ API密钥认证成功");
            println!("    ✅ 用户: {}", user.username);
            println!("    ✅ 角色: {:?}", user.roles);
            
            // 测试权限检查
            println!("  🛡️ 测试权限检查...");
            let authz_result = auth_provider.authorize(&user, "/api/users", "read").await;
            match authz_result {
                data_gateway_gateway::middleware::AuthzResult::Allowed => {
                    println!("    ✅ 权限检查通过");
                }
                data_gateway_gateway::middleware::AuthzResult::Denied(reason) => {
                    println!("    ❌ 权限检查失败: {}", reason);
                }
            }
        }
        data_gateway_gateway::middleware::AuthResult::Failed(reason) => {
            println!("    ❌ API密钥认证失败: {}", reason);
        }
        data_gateway_gateway::middleware::AuthResult::Required => {
            println!("    ⚠️ 需要认证");
        }
    }

    println!("  🔐 测试基础认证...");
    let basic_auth_method = AuthMethod::BasicAuth {
        username: "admin".to_string(),
        password: "admin123".to_string(),
    };
    let auth_result = auth_provider.authenticate(basic_auth_method).await;
    match auth_result {
        data_gateway_gateway::middleware::AuthResult::Success(user) => {
            println!("    ✅ 基础认证成功");
            println!("    ✅ 用户: {}", user.username);
        }
        data_gateway_gateway::middleware::AuthResult::Failed(reason) => {
            println!("    ❌ 基础认证失败: {}", reason);
        }
        data_gateway_gateway::middleware::AuthResult::Required => {
            println!("    ⚠️ 需要认证");
        }
    }

    println!("  🔍 测试认证信息提取...");
    let mut headers = axum::http::HeaderMap::new();
    headers.insert("X-API-Key", "test-key".parse().unwrap());
    
    let auth_method = MiddlewareManager::extract_auth_method(&headers);
    match auth_method {
        Some(AuthMethod::ApiKey(key)) => {
            println!("    ✅ 成功提取API密钥: {}", key);
        }
        _ => {
            println!("    ❌ 未能提取认证信息");
        }
    }

    Ok(())
}

/// 演示S3对象存储数据源功能
async fn demo_s3_data_source() -> anyhow::Result<()> {
    println!("  创建S3数据源配置...");
    let s3_config = S3Config {
        endpoint: "https://s3.amazonaws.com".to_string(),
        access_key_id: "demo_access_key".to_string(),
        secret_access_key: "demo_secret_key".to_string(),
        region: "us-east-1".to_string(),
        bucket: "demo-bucket".to_string(),
        prefix: Some("data/".to_string()),
    };

    println!("  🔗 创建S3数据源...");
    let mut s3_source = S3DataSource::new(s3_config);
    println!("    ✅ S3数据源创建成功");
    println!("    ✅ 数据源名称: {}", s3_source.name());
    println!("    ✅ 数据源描述: {}", s3_source.description());

    println!("  🔌 连接到S3...");
    s3_source.connect().await?;
    println!("    ✅ S3连接成功");

    println!("  🏥 健康检查...");
    let health = s3_source.health_check().await?;
    println!("    ✅ 健康检查结果: {}", health);

    println!("  📋 执行查询（列出对象）...");
    let query = DataSourceQuery {
        id: "list_objects_query".to_string(),
        content: "SELECT * FROM objects".to_string(),
        parameters: HashMap::new(),
    };
    
    let result = s3_source.execute_query(query).await?;
    if result.success {
        println!("    ✅ 查询执行成功");
        println!("    ✅ 处理行数: {}", result.stats.rows_processed);
        println!("    ✅ 执行时间: {:.2}ms", result.stats.execution_time_ms);
        println!("    ✅ 返回数据: {}", result.data);
    } else {
        println!("    ❌ 查询执行失败: {:?}", result.error);
    }

    println!("  🔌 断开S3连接...");
    s3_source.disconnect().await?;
    println!("    ✅ S3连接断开成功");

    Ok(())
}

/// 演示Elasticsearch搜索引擎数据源功能
async fn demo_elasticsearch_data_source() -> anyhow::Result<()> {
    println!("  创建Elasticsearch数据源配置...");
    let es_config = ElasticsearchConfig {
        url: "http://localhost:9200".to_string(),
        username: Some("elastic".to_string()),
        password: Some("password".to_string()),
        default_index: "logs".to_string(),
        connect_timeout: 30,
        query_timeout: 60,
    };

    println!("  🔗 创建Elasticsearch数据源...");
    let mut es_source = ElasticsearchDataSource::new(es_config);
    println!("    ✅ Elasticsearch数据源创建成功");
    println!("    ✅ 数据源名称: {}", es_source.name());
    println!("    ✅ 数据源描述: {}", es_source.description());

    println!("  🏥 健康检查...");
    let health = es_source.health_check().await?;
    println!("    ✅ 健康检查结果: {}", health);

    println!("  📋 SQL到ES查询转换演示...");
    let query = DataSourceQuery {
        id: "search_logs_query".to_string(),
        content: "SELECT * FROM logs WHERE level = 'ERROR'".to_string(),
        parameters: HashMap::new(),
    };

    let result = es_source.execute_query(query).await?;
    if result.success {
        println!("    ✅ 查询转换成功");
        println!("    ✅ 执行时间: {:.2}ms", result.stats.execution_time_ms);
    } else {
        println!("    ⚠️ 查询执行失败（预期，因为未连接到真实ES集群）: {:?}", result.error);
    }

    println!("  🔍 测试不同类型的查询...");
    let queries = vec![
        "SELECT * FROM logs",
        "SELECT timestamp, message FROM logs WHERE level = 'INFO'",
        "SELECT COUNT(*) FROM logs GROUP BY level",
    ];

    for (i, sql) in queries.iter().enumerate() {
        println!("    📝 查询 {}: {}", i + 1, sql);
        let query = DataSourceQuery {
            id: format!("test_query_{}", i + 1),
            content: sql.to_string(),
            parameters: HashMap::new(),
        };

        let result = es_source.execute_query(query).await?;
        println!("    ✅ 查询转换完成，执行时间: {:.2}ms", result.stats.execution_time_ms);
    }

    Ok(())
}

/// 演示限流控制中间件功能
async fn demo_rate_limiting() -> anyhow::Result<()> {
    println!("  创建限流配置...");
    let rate_limit_config = RateLimitConfig {
        requests_per_minute: 10,
        burst_capacity: 5,
        window_size: 60,
    };

    println!("  🚦 创建中间件管理器...");
    let mut middleware_config = data_gateway_gateway::middleware::MiddlewareConfig::default();
    middleware_config.enable_rate_limit = true;
    middleware_config.rate_limit_config = Some(rate_limit_config);

    let manager = MiddlewareManager::new(middleware_config);
    println!("    ✅ 中间件管理器创建成功");

    println!("  🔄 测试限流功能...");
    let client_id = "test_client_123";

    // 模拟多次请求
    for i in 1..=12 {
        let allowed = manager.check_rate_limit(client_id).await.unwrap_or(false);
        let remaining = manager.get_remaining_requests(client_id).await;

        if allowed {
            println!("    ✅ 请求 {} 被允许，剩余请求数: {}", i, remaining);
        } else {
            println!("    ❌ 请求 {} 被限流，剩余请求数: {}", i, remaining);
        }
    }

    println!("  🧹 测试限流记录清理...");
    manager.cleanup_rate_limit_records().await;
    println!("    ✅ 过期记录清理完成");

    println!("  🌐 测试客户端ID提取...");
    let mut headers = axum::http::HeaderMap::new();
    headers.insert("X-Forwarded-For", "192.168.1.100, 10.0.0.1".parse().unwrap());

    let client_id = MiddlewareManager::extract_client_id(&headers, Some("127.0.0.1"));
    println!("    ✅ 提取的客户端ID: {}", client_id);

    // 测试不同的客户端
    println!("  👥 测试多客户端限流...");
    let clients = vec!["client_a", "client_b", "client_c"];

    for client in &clients {
        for i in 1..=3 {
            let allowed = manager.check_rate_limit(client).await.unwrap_or(false);
            println!("    {} 请求 {}: {}", client, i, if allowed { "✅ 允许" } else { "❌ 限流" });
        }
    }

    Ok(())
}

/// 演示综合功能
async fn demo_integrated_features() -> anyhow::Result<()> {
    println!("  🎯 综合功能演示：多组件协同工作");

    // 创建数据源管理器并注册多个数据源
    println!("  📦 创建数据源管理器...");
    let mut ds_manager = DataSourceManager::new();

    // 注册S3数据源
    let s3_config = S3Config::default();
    let s3_source = Box::new(S3DataSource::new(s3_config));
    ds_manager.register_source("s3".to_string(), s3_source);

    // 注册Elasticsearch数据源
    let es_config = ElasticsearchConfig::default();
    let es_source = Box::new(ElasticsearchDataSource::new(es_config));
    ds_manager.register_source("elasticsearch".to_string(), es_source);

    println!("    ✅ S3和Elasticsearch数据源注册成功");

    // 列出所有数据源
    let sources = ds_manager.list_sources();
    println!("    ✅ 已注册数据源: {:?}", sources);

    // 创建Flight SQL服务
    println!("  🛩️ 创建Flight SQL服务...");
    let kernel = Arc::new(DataGatewayKernel::new().await?);
    let flight_service = SimpleFlightSqlService::new(kernel);

    // 创建限流中间件
    println!("  🚦 创建限流中间件...");
    let rate_limit_config = RateLimitConfig {
        requests_per_minute: 100,
        burst_capacity: 20,
        window_size: 60,
    };
    let mut middleware_config = data_gateway_gateway::middleware::MiddlewareConfig::default();
    middleware_config.enable_rate_limit = true;
    middleware_config.rate_limit_config = Some(rate_limit_config);
    let middleware_manager = MiddlewareManager::new(middleware_config);

    // 在事务中执行准备语句
    println!("  🔄 在事务中执行准备语句...");
    let tx_id = flight_service.begin_transaction().await?;
    let stmt_id = flight_service.create_prepared_statement("SELECT COUNT(*) FROM multi_source_data").await?;

    println!("    ✅ 事务ID: {}", tx_id);
    println!("    ✅ 准备语句ID: {}", stmt_id);

    // 测试限流功能
    println!("  🚦 测试限流功能...");
    for i in 1..=5 {
        let allowed = middleware_manager.check_rate_limit("integrated_test_client").await.unwrap_or(false);
        println!("    请求 {}: {}", i, if allowed { "✅ 允许" } else { "❌ 限流" });
    }

    // 提交事务
    flight_service.commit_transaction(&tx_id).await?;
    flight_service.close_prepared_statement(&stmt_id).await?;

    println!("    ✅ 事务提交并清理完成");

    // 演示多数据源查询协调
    println!("  🔗 演示多数据源查询协调...");
    println!("    📊 S3数据源：处理对象存储查询");
    println!("    🔍 Elasticsearch数据源：处理全文搜索查询");
    println!("    🛩️ Flight SQL：提供统一的SQL接口");
    println!("    🚦 限流中间件：保护系统资源");

    println!("  🎉 综合功能演示完成！所有组件协同工作正常");
    println!("    ✅ 数据源插件：多种存储后端支持");
    println!("    ✅ Flight SQL：企业级SQL接口");
    println!("    ✅ 中间件：安全和性能保护");
    println!("    ✅ 微内核架构：高度可扩展");

    Ok(())
}
